package com.bridgeintelligent.tag.engine.es.base;

import com.bridgeintelligent.tag.engine.api.EngineException;
import com.bridgeintelligent.tag.engine.es.index.EngineMappingHelper;
import com.bridgeintelligent.tag.engine.es.model.CustomerGroupModel;
import com.bridgeintelligent.tag.engine.es.model.DataSetModel;
import com.bridgeintelligent.tag.engine.es.model.IndexModel;
import com.wayneleo.quickstart.framework.ExceptionCode;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.*;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.indices.GetMappingsRequest;
import org.elasticsearch.client.indices.GetMappingsResponse;
import org.elasticsearch.cluster.metadata.MappingMetadata;
import org.elasticsearch.common.Strings;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * Created by JackCheung on 2019/9/24 5:10 PM.
 */
@Slf4j
public class EngineESHelper implements IEngineESHelper {

    private RestHighLevelClient client;

    /**
     * 配置es客户端
     *
     * @param esConfig EngineESConfig
     */
    EngineESHelper(EngineESConfig esConfig) {
        String[] address = esConfig.getAddress().split(";");
        HttpHost[] httpHosts = new HttpHost[address.length];
        for (int i = 0; i < address.length; i++) {
            httpHosts[i] = new HttpHost(StringUtils.substringBefore(address[i], ":"),
                                        Integer.parseInt(StringUtils.substringAfter(address[i], ":")),
                                        "http");
        }
        final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(AuthScope.ANY,
                                           new UsernamePasswordCredentials(esConfig.getUsername(),
                                                                           esConfig.getPassword()));
        client = new RestHighLevelClient(RestClient.builder(httpHosts)
                                                   .setHttpClientConfigCallback(httpAsyncClientBuilder -> httpAsyncClientBuilder
                                                           .setDefaultCredentialsProvider(credentialsProvider)
                                                           .setKeepAliveStrategy((httpResponse, httpContext) -> TimeUnit.MINUTES
                                                                   .toMinutes(6)))
                                                    .setRequestConfigCallback(requestConfigCallback -> requestConfigCallback
                                                            .setConnectTimeout(600000)
                                                            .setSocketTimeout(600000)));
    }

    @Override
    public RestHighLevelClient getClient() {
        return client;
    }

    /**
     * 创建索引字段映射
     *
     * @param indexName indexName
     */
    @Override
    public void updateIndexMapping(String indexName) {
        PutMappingRequest request = new PutMappingRequest(indexName);
        request.type(EngineMappingHelper.TYPE);
        request.source(EngineMappingHelper.updateMapping());
        try {
            client.indices().putMapping(request,RequestOptions.DEFAULT);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @Override
    public Set<String> getIndexMapping(String indexName)  {
        GetMappingsRequest request = new GetMappingsRequest().indices(indexName);
        GetMappingsResponse response;
        try {
            response = client.indices().getMapping(request, RequestOptions.DEFAULT);
            Map<String, MappingMetadata> mappings = response.mappings();
            MappingMetadata indexMapping = mappings.get(indexName);
            Map<String, Object> mapping = indexMapping.sourceAsMap();
            Map<String,Object> properties = (Map<String,Object> )mapping.get("properties");
            return properties.keySet();
        }catch (IOException e){
            e.printStackTrace();
            return new HashSet<>();
        }
    }

    /**
     * 批量新增
     *
     * @param indexName indexName
     * @param map       map
     */
    @Override
    public void bulkInsert(String indexName, Map<String, Map<String, Object>> map) {
        BulkRequest bulkRequest = new BulkRequest();
        map.forEach((k, v) -> bulkRequest.add(new IndexRequest(indexName, EngineMappingHelper.TYPE, k).source(v)));
        try {
            BulkResponse responses = client.bulk(bulkRequest, RequestOptions.DEFAULT);
            responses.hasFailures();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 批量更新
     *
     * @param indexName indexName
     * @param map       map
     */
    @Override
    public void bulkUpdate(String indexName, Map<String, Map<String, Object>> map) {
        BulkRequest bulkRequest = new BulkRequest();
        map.forEach((k, v) -> bulkRequest.add(new UpdateRequest(indexName, EngineMappingHelper.TYPE, k).doc(v)));
        try {
            client.bulk(bulkRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 根据索引名和查询条件查询数据
     *
     * @param indexName    indexName
     * @param dataSetModel dataSetModel
     * @return SearchResponse
     */
    @Override
    public SearchResponse dataSetSearch(String indexName, DataSetModel dataSetModel) {

        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        if (dataSetModel.getSortBuilder() != null) {
            searchSourceBuilder.sort(dataSetModel.getSortBuilder());
        }
        searchSourceBuilder.query(dataSetModel.getQueryBuilder())
                           .from(dataSetModel.getFrom())
                           .size(dataSetModel.getSize())
                           .trackTotalHits(true)
                           .timeout(new TimeValue(60, TimeUnit.SECONDS));
        if (dataSetModel.getIncludes() != null && dataSetModel.getIncludes().length != 0) {
            searchSourceBuilder.fetchSource(dataSetModel.getIncludes(), dataSetModel.getExcludes());
        }
        log.info("dataSet查询条件为----->>>>>{}", searchSourceBuilder);
        //创建SearchRequest
        SearchRequest searchRequest = new SearchRequest(indexName).source(searchSourceBuilder);
        //查询结果并返回
        return createResult(searchRequest, "dataSet");
    }

    @Override
    public SearchResponse querySearch(String indexName) {
        return null;
    }

    /**
     * 根据indexName和customerGroupModel查询结果
     *
     * @param indexName          indexName
     * @param customerGroupModel customerGroupModel
     * @return SearchResponse
     */
    @Override
    public SearchResponse  customerGroup(String indexName, CustomerGroupModel customerGroupModel) {
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.from(0)
                           .size(0)
                           .query(customerGroupModel.getQueryBuilder())
                           .trackTotalHits(true)
                           .timeout(new TimeValue(60, TimeUnit.SECONDS));
        customerGroupModel.getAggregationBuilderList().forEach(searchSourceBuilder::aggregation);
        log.info("\ncustomerGroup查询条件为----->>>>>{}", searchSourceBuilder);
        SearchRequest searchRequest = new SearchRequest(indexName).source(searchSourceBuilder);
        return createResult(searchRequest, "customerGroup");
    }

    /**
     * 聚合分析
     *
     * @param indexModel indexModel
     * @param indexName  indexName
     * @return List<SearchResponse>
     */
    @Override
    public List<SearchResponse> indexSearch(IndexModel indexModel, String... indexName) {
        //创建searchSourceBuilder
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.from(0)
                           .size(0)
                           .query(indexModel.getQueryBuilder())
                           .trackTotalHits(true)
                           .timeout(new TimeValue(60, TimeUnit.SECONDS));
        if (null != indexModel.getAggregationBuilder()) {
            //添加聚合
            searchSourceBuilder.aggregation(indexModel.getAggregationBuilder());
        }
        log.info("\nindex查询条件为----->>>>>{}", searchSourceBuilder);
        List<SearchResponse> result = new ArrayList<>(indexName.length);
        for (String indexNameItem : indexName) {
            SearchRequest searchRequest = new SearchRequest(indexNameItem).source(searchSourceBuilder);
            result.add(createResult(searchRequest, "indexSearch"));
        }
        return result;
    }

    @Override
    public List<SearchHit[]> scrollSearch(String indexName, BoolQueryBuilder boolQueryBuilder, Set<String> fieldsSet) {
        List<SearchHit[]> searchHitList = new ArrayList<>();
        Scroll scroll = new Scroll(TimeValue.timeValueMinutes(1L));
        //封装查询条件
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        //封装需要查询出来的字段
        String[] includes = new String[fieldsSet.size()];
        searchSourceBuilder.query(boolQueryBuilder)
                           .trackTotalHits(true)
                           .fetchSource(fieldsSet.toArray(includes), Strings.EMPTY_ARRAY)
                           .timeout(new TimeValue(60, TimeUnit.SECONDS));
        //封装请求
        SearchRequest searchRequest = new SearchRequest(indexName).scroll(scroll).source(searchSourceBuilder);
        try {
            SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
            String scrollId = searchResponse.getScrollId();
            SearchHit[] searchHits = searchResponse.getHits().getHits();
            while (searchHits != null && searchHits.length > 0) {
                searchHitList.add(searchHits);
                SearchScrollRequest scrollRequest = new SearchScrollRequest(scrollId);
                scrollRequest.scroll(scroll);
                searchResponse = client.searchScroll(scrollRequest, RequestOptions.DEFAULT);
                scrollId = searchResponse.getScrollId();
                searchHits = searchResponse.getHits().getHits();
            }
            ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
            clearScrollRequest.addScrollId(scrollId);
            ClearScrollResponse clearScrollResponse = client.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
            boolean succeeded = clearScrollResponse.isSucceeded();
            log.info("------清理上下文结果------" + succeeded);
        } catch (IOException e) {
            log.error("dataSet-----scrollSearch查询失败", e);
        }
        return searchHitList;
    }

    /**
     * 从ES中查询结果并返回SearchResponse
     *
     * @param searchRequest searchRequest
     * @param scene         scene
     * @return SearchResponse
     */
    private SearchResponse createResult(SearchRequest searchRequest, String scene) {
        try {
            //从ES中查询结果
            SearchResponse result = client.search(searchRequest, RequestOptions.DEFAULT);
            //log.info("\n{}查询返回结果-------->>>{}", scene, result);
            return result;
        } catch (ElasticsearchStatusException e) {
            log.error("查询发生ElasticsearchStatusException异常--------->>>>", e);
            throw new EngineException(new ExceptionCode(999, e.getMessage()));
        } catch (IOException e) {
            e.printStackTrace();
            log.error("{}查询发生异常-------->>>{}", scene, e.getMessage());
            throw new EngineException(new ExceptionCode(999, e.getMessage()));
        }
    }
}
